/*
 * SPDX-FileCopyrightText: 2020-2025 Espressif Systems (Shanghai) CO LTD
 * SPDX-License-Identifier: Apache-2.0
 */
#include "sdkconfig.h"

#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_chip_info.h"
#include "usbd_core.h"
#include "esp_jtag_main.h"

#include "esp_jtag_vendord.h"
#include "esp_io.h"

static const char *TAG = "esp_usb_jtag"; // 日志标签

// 全局变量定义
static esp_chip_model_t s_target_model;        // 存储目标芯片型号(ESP32/ESP32-S2等)
static TaskHandle_t s_jtag_task_handle = NULL; // JTAG任务句柄，用于任务挂起/恢复控制
static uint8_t s_tdo_bytes[1024];              // TDO(测试数据输出)数据缓冲区(1024字节)
static uint16_t s_total_tdo_bits = 0;          // 已接收的TDO总位数
static uint16_t s_usb_sent_bits = 0;           // 已通过USB发送的TDO位数

// 工具宏：将位数向上对齐到字节边界
#define ROUND_UP_BITS(x) ((x + 7) & (~7))

/* JTAG协议相关定义 */
#define JTAG_PROTO_MAX_BITS (EUB_VENDORD_EPSIZE * 8) // 单次传输最大位数(基于USB端点大小)
#define JTAG_PROTO_CAPS_VER 1                        // 协议版本号

// JTAG能力描述头结构体
typedef struct __attribute__((packed)) {
    uint8_t proto_ver; // 协议版本(必须为JTAG_PROTO_CAPS_VER)
    uint8_t length;    // 后续描述符总长度
} jtag_proto_caps_hdr_t;

// APB速度类型标识
#define JTAG_PROTO_CAPS_SPEED_APB_TYPE 1

// 通用描述头结构体
typedef struct __attribute__((packed)) {
    uint8_t type;   // 描述符类型
    uint8_t length; // 描述符长度
} jtag_gen_hdr_t;

// APB速度描述结构体
typedef struct __attribute__((packed)) {
    uint8_t type;             // 必须为JTAG_PROTO_CAPS_SPEED_APB_TYPE
    uint8_t length;           // 结构体长度
    uint16_t apb_speed_10khz; // APB总线速度(单位：10KHz)
    uint16_t div_min;         // 最小分频系数
    uint16_t div_max;         // 最大分频系数
} jtag_proto_caps_speed_apb_t;

// 完整的JTAG能力描述结构体
typedef struct
{
    jtag_proto_caps_hdr_t proto_hdr;      // 协议头
    jtag_proto_caps_speed_apb_t caps_apb; // APB速度描述
} jtag_proto_caps_t;

/* 厂商自定义请求码 */
#define VEND_JTAG_SETDIV     0 // 设置时钟分频
#define VEND_JTAG_SETIO      1 // 设置IO状态
#define VEND_JTAG_GETTDO     2 // 获取TDO数据
#define VEND_JTAG_SET_CHIPID 3 // 设置目标芯片ID

// 默认TCK频率定义(4800KHz，实际为4800*2/10=960KHz)
#define TCK_FREQ(khz) ((khz * 2) / 10)

// 默认JTAG能力描述(固定参数)
static const jtag_proto_caps_t jtag_proto_caps = {
    { .proto_ver = JTAG_PROTO_CAPS_VER, .length = sizeof(jtag_proto_caps_hdr_t) + sizeof(jtag_proto_caps_speed_apb_t) },
    { .type = JTAG_PROTO_CAPS_SPEED_APB_TYPE, .length = sizeof(jtag_proto_caps_speed_apb_t), .apb_speed_10khz = TCK_FREQ(4800), .div_min = 1, .div_max = 1 }
};

/* 函数别名定义(兼容旧版API) */
void eub_debug_probe_init(void) __attribute__((alias("esp_usb_jtag_init")));
int eub_debug_probe_get_proto_caps(void *dest) __attribute__((alias("esp_usb_jtag_get_proto_caps")));
void eub_debug_probe_task_suspend(void) __attribute__((alias("esp_usb_jtag_task_suspend")));
void eub_debug_probe_task_resume(void) __attribute__((alias("esp_usb_jtag_task_resume")));
bool eub_debug_probe_target_is_esp32(void) __attribute__((alias("esp_usb_jtag_target_is_esp32")));

/*                 在获取字符串描述符时调用，修改Cherryusb源码                 */
/* 获取JTAG协议能力描述 */
static int esp_usb_jtag_get_proto_caps(void *dest)
{
    memcpy(dest, (uint16_t *)&jtag_proto_caps, sizeof(jtag_proto_caps));
    return sizeof(jtag_proto_caps);
}

/* 检查目标芯片是否为ESP32 */
static bool esp_usb_jtag_target_is_esp32(void)
{
    return s_target_model == CHIP_ESP32;
}

/* 挂起JTAG任务 */
static void esp_usb_jtag_task_suspend(void)
{
    if (s_jtag_task_handle) {
        vTaskSuspend(s_jtag_task_handle); // 挂起FreeRTOS任务
    }
}

/* 恢复JTAG任务 */
static void esp_usb_jtag_task_resume(void)
{
    if (s_jtag_task_handle) {
        vTaskResume(s_jtag_task_handle); // 恢复FreeRTOS任务
    }
}

/* USB厂商请求处理函数 */
int esp_usb_jtag_debug_probe_control_handler(uint8_t busid, struct usb_setup_packet *setup)
{
    switch (setup->bRequest) {
        case VEND_JTAG_SETDIV: // 设置时钟分频(暂未实现)
        case VEND_JTAG_SETIO:  // 设置IO状态(暂未实现)
            //return tud_control_status(rhport, request);
            return 0;
        case VEND_JTAG_GETTDO: {                    // 获取TDO状态
            uint8_t buf = gpio_get_level(GPIO_TDO); // 读取TDO引脚电平
            return usbd_ep_start_write(0, ESP_JTAG_IN_EP, (void *)&buf, 1);
            //return tud_control_xfer(rhport, request, (void *)&buf, 1);
        }

        case VEND_JTAG_SET_CHIPID: // 设置目标芯片ID
            s_target_model = setup->wValue;
            return 0;
            //return tud_control_status(rhport, request);

        default:
            return -1; // 未知请求
    }
}

/* 执行单次JTAG操作(时钟周期) */
inline static void do_jtag_one(const uint8_t tdo_req, const uint8_t tms_tdi_mask)
{
    esp_gpio_write_tmstck(tms_tdi_mask); // 设置TMS/TDI状态
    esp_gpio_tck_set();                  // TCK上升沿

    if (tdo_req) { // 如果需要读取TDO
        // 将TDO数据存入缓冲区，按位存储
        s_tdo_bytes[s_total_tdo_bits / 8] |= (esp_gpio_tdo_read() << (s_total_tdo_bits % 8));
        s_total_tdo_bits++;
    }

    esp_gpio_tck_clr(); // TCK下降沿
}

/* JTAG主任务 */
static void esp_usb_jtag_task(void *pvParameters)
{
    // JTAG命令枚举
    enum e_cmds {
        CMD_CLK_0 = 0, // TMS=0, TDI=0
        CMD_CLK_1,     // TMS=0, TDI=1
        CMD_CLK_2,     // TMS=1, TDI=0
        CMD_CLK_3,     // TMS=1, TDI=1
        CMD_CLK_4,     // 读TDO, TMS=0, TDI=0
        CMD_CLK_5,     // 读TDO, TMS=0, TDI=1
        CMD_CLK_6,     // 读TDO, TMS=1, TDI=0
        CMD_CLK_7,     // 读TDO, TMS=1, TDI=1
        CMD_SRST0,     // TAP复位(TMS=1)
        CMD_SRST1,     // 系统复位(保留)
        CMD_FLUSH,     // 刷新TDO数据
        CMD_RSV,       // 保留命令
        CMD_REP0,      // 重复命令0
        CMD_REP1,      // 重复命令1
        CMD_REP2,      // 重复命令2
        CMD_REP3       // 重复命令3
    };

    // 命令对应的TDO请求和TMS/TDI状态
    const struct
    {
        uint8_t tdo_req;      // 是否读取TDO
        uint8_t tms_tdi_mask; // TMS/TDI掩码
    } pin_levels[] = {
        { 0, 0 },                 // CMD_CLK_0: TMS=0, TDI=0
        { 0, GPIO_TDI_MASK },     // CMD_CLK_1: TMS=0, TDI=1
        { 0, GPIO_TMS_MASK },     // CMD_CLK_2: TMS=1, TDI=0
        { 0, GPIO_TMS_TDI_MASK }, // CMD_CLK_3: TMS=1, TDI=1
        { 1, 0 },                 // CMD_CLK_4: 读TDO, TMS=0, TDI=0
        { 1, GPIO_TDI_MASK },     // CMD_CLK_5: 读TDO, TMS=0, TDI=1
        { 1, GPIO_TMS_MASK },     // CMD_CLK_6: 读TDO, TMS=1, TDI=0
        { 1, GPIO_TMS_TDI_MASK }, // CMD_CLK_7: 读TDO, TMS=1, TDI=1
        { 0, GPIO_TMS_MASK },     // CMD_SRST0: TAP复位(TMS=1)
        { 0, GPIO_TMS_MASK },     // CMD_SRST1: 系统复位(保留)
    };

    s_jtag_task_handle = xTaskGetCurrentTaskHandle(); // 保存当前任务句柄

    size_t cnt = 0;
    int prev_cmd = CMD_SRST0, rep_cnt = 0; // 前一个命令和重复计数器

    while (1) {
        //esp_gpio_jtag_led_off();                             // LED指示开始处理
        char *nibbles = eub_vendord_recv_acquire_item(&cnt);   // 从USB获取数据(半字节数组)
        //esp_gpio_jtag_led_on();                              // LED指示处理中
        // ESP_LOGI(TAG, "Received %d nibbles", cnt);
        // for(int i = 0; i < cnt; i++)
        //     ESP_LOGI(TAG, "nibbles[%d] = %d",cnt, *nibbles);

        ESP_LOG_BUFFER_HEXDUMP(TAG, nibbles, cnt, ESP_LOG_DEBUG); // 调试日志

        // 处理每个半字节(4bit)
        for (size_t n = 0; n < cnt * 2; n++) {
            // 从半字节中提取命令(4bit)
            const int cmd = (n & 1) ? (nibbles[n / 2] & 0x0F) : (nibbles[n / 2] >> 4);

            int cmd_exec = cmd, cmd_rpt_cnt = 1; // 实际执行的命令和重复次数

            // 命令解析
            switch (cmd) {
                case CMD_REP0 ... CMD_REP3: // 重复命令
                    cmd_rpt_cnt = (cmd - CMD_REP0) << (2 * rep_cnt++);
                    cmd_exec = prev_cmd;
                    break;
                case CMD_SRST0: // TAP复位(强制发送8个TMS=1)
                    cmd_rpt_cnt = 8;
                    break;
                case CMD_SRST1: // 系统复位(保留实现)
                    cmd_rpt_cnt = 8;
                    break;
                default:
                    rep_cnt = 0; // 重置重复计数器
                    break;
            }

            // 执行JTAG操作
            if (cmd_exec < CMD_FLUSH) {
                for (int i = 0; i < cmd_rpt_cnt; i++) {
                    do_jtag_one(pin_levels[cmd_exec].tdo_req, pin_levels[cmd_exec].tms_tdi_mask);
                }
            }
            // 刷新TDO数据到USB
            else if (cmd_exec == CMD_FLUSH) {
                s_total_tdo_bits = ROUND_UP_BITS(s_total_tdo_bits); // 对齐到字节边界
                if (s_usb_sent_bits < s_total_tdo_bits) {
                    int waiting_to_send_bits = s_total_tdo_bits - s_usb_sent_bits;
                    while (waiting_to_send_bits > 0) {
                        // 计算本次要发送的位数(不超过最大限制)
                        int send_bits = MIN(waiting_to_send_bits, JTAG_PROTO_MAX_BITS);
                        eub_vendord_send_item(s_tdo_bytes + (s_usb_sent_bits / 8), send_bits / 8);
                        s_usb_sent_bits += send_bits;
                        waiting_to_send_bits -= send_bits;
                    }
                    memset(s_tdo_bytes, 0x00, sizeof(s_tdo_bytes)); // 清空缓冲区
                    s_total_tdo_bits = s_usb_sent_bits = 0;
                }
            }

            // 如果累积数据达到阈值，立即发送
            int waiting_to_send_bits = s_total_tdo_bits - s_usb_sent_bits;
            if (waiting_to_send_bits >= JTAG_PROTO_MAX_BITS) {
                int send_bits = ROUND_UP_BITS(MIN(waiting_to_send_bits, JTAG_PROTO_MAX_BITS));
                int n_byte = send_bits / 8;
                eub_vendord_send_item(s_tdo_bytes + (s_usb_sent_bits / 8), n_byte);
                memset(s_tdo_bytes + (s_usb_sent_bits / 8), 0x00, n_byte);
                s_usb_sent_bits += send_bits;
                waiting_to_send_bits -= send_bits;
                if (waiting_to_send_bits <= 0) {
                    s_total_tdo_bits = s_usb_sent_bits = 0;
                }
            }

            // 更新前一个命令(非重复/刷新命令时)
            if (cmd < CMD_REP0 && cmd != CMD_FLUSH) {
                prev_cmd = cmd;
            }
        }

        eub_vendord_free_item(nibbles); // 释放数据缓冲区
    }

    vTaskDelete(NULL); // 理论上不会执行到这里
}
/* JTAG初始化函数 */
static void esp_usb_jtag_init(void)
{
    esp_init_jtag_pins(); // 初始化JTAG引脚

    // 创建JTAG任务(绑定到当前核心)
    BaseType_t res = xTaskCreatePinnedToCore(
        esp_usb_jtag_task,            // 任务函数
        "jtag_task",                  // 任务名称
        4 * 1024,                     // 堆栈大小(4KB)
        NULL,                         // 参数
        EUB_VENDORD_USB_TASK_PRI + 1, // 优先级(高于USB任务)
        NULL,                         // 任务句柄
        esp_cpu_get_core_id()         // 绑定到当前核心
    );

    if (res != pdPASS) {
        ESP_LOGE(TAG, "Cannot create ESP_JTAG task!");
        abort(); // 创建失败时中止
    }
}
/* JTAG反初始化函数 */
void esp_usb_jtag_deinit(void)
{
    esp_reset_dap_pins(); // 复位JTAG引脚
    if (s_jtag_task_handle){
        vTaskDelete(s_jtag_task_handle); // 删除JTAG任务
        s_jtag_task_handle = NULL;
    }
}   
